package com.easy.lab.shixun.controller.ws;

import com.easy.lab.shixun.api.service.GroupMemberService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;

import java.security.Principal;

/**
 * 监控ws上下线自定义处理
 * 本例子缓存了ws用户令牌和ws会话关系
 * @Author: https://gitee.com/wesleyOne
 * @Date: 12.07 2019
 */
@Component
@Slf4j
public class SimpleWebSocketHandlerDecoratorFactory implements WebSocketHandlerDecoratorFactory {

    @Autowired
    private GroupMemberService groupMemberService;

    @Override
    public WebSocketHandler decorate(WebSocketHandler handler) {
        return new WebSocketHandlerDecorator(handler) {

            @Override
            public void afterConnectionEstablished(WebSocketSession session) throws Exception {
                log.debug("建立会话  sessionId = {}", session.getId());
                Principal principal = session.getPrincipal();
                if (principal != null) {
                    log.info("关联token和会话 token: {}", principal.getName());
                    groupMemberService.relateTokenWithWebSocketSession(principal.getName(),session);
                }
                super.afterConnectionEstablished(session);
            }

            @Override
            public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
                log.debug("会话断开  sessionId = {}", session.getId());
                Principal principal = session.getPrincipal();
                if (principal != null) {
                    // 身份校验成功，移除socket连接
                    log.info("移除token和会话关联 token: {}",principal.getName());
                    groupMemberService.removeTokenWsSession(principal.getName());
                }
                super.afterConnectionClosed(session, closeStatus);
            }
        };
    }
}
